package mike.util.csv;

import java.io.*;
import java.sql.*;
import java.util.*;

public class CSVWriter implements Closeable
{
	/**
	 * The default field separator character (comma)
	 */
	public static final char DEFAULT_SEPARATOR = CSVReader.DEFAULT_SEPARATOR;
	
	/**
	 * The underlying file writer.
	 */
	private BufferedWriter writer;
	/**
	 * The separator character used by this instance.
	 */
	private char separator;
	/**
	 * The separator held as a String for use by the encode method
	 */
	private String separatorString;
	/**
	 * If true, each line that is written will have the same number of fields, which is based
	 * off the number of fields in the header row (which must be supplied if this is true).
	 * If a line that is to be written has fewer fields, then empty strings are added
	 * on to the end of the row.  If a line that is to be written has more fields, then
	 * the extra fields are ignored.
	 */
	private boolean keepNumberOfColumnsConsistent;
	/**
	 * The number of columns to be written if keepNumberOfColumnsConsistent is true.
	 */
	private int numCols;
	/**
	 * The number of columns written in the current row (when writing by fields instead of rows).
	 * Used to restrict the number of columns if keepNumberOfColumnsConsistent is true.
	 */
	private int numColsThisRow;
	/**
	 * Whether or not we are about to write a new row, ie, the last thing
	 * to be written was a newline.
	 */
	private boolean newRow = true;
	
	/**
	 * Create a new writer instance with the specified filename and default separator
	 * @param filename The filename of the file to be written
	 * @throws IOException If the file cannot be created or written
	 * @throws NullPointerException If filename is null
	 */
	public CSVWriter(String filename) throws IOException
	{
		this(filename, DEFAULT_SEPARATOR);
	}
	/**
	 * Create a new writer instance with the specified filename and separator
	 * @param filename The filename of the file to be written
	 * @param separator The character to be used as the field separator
	 * @throws IOException If the file cannot be created or written
	 * @throws NullPointerException If filename is null
	 */
	public CSVWriter(String filename, char separator) throws IOException
	{
		constructorCommon(filename, separator);
	}
	/**
	 * Create a new writer instance with the specified filename and header values
	 * using the default separator.
	 * @param filename The filename of the file to be written
	 * @param headers The values of each field in the header row
	 * @throws IOException If the file cannot be created or written
	 * @throws NullPointerException If filename is null
	 * @throws NullPointerException If headers is null
	 */
	public CSVWriter(String filename, String[] headers) throws IOException
	{
		this(filename, DEFAULT_SEPARATOR, headers, false);
	}
	/**
	 * Create a new writer instance with the specified filename, separator, and header values
	 * @param filename The filename of the file to be written
	 * @param separator The character to be used as the field separator
	 * @throws IOException If the file cannot be created or written
	 * @throws NullPointerException If filename is null
	 * @throws NullPointerException If headers is null
	 */
	public CSVWriter(String filename, char separator, String[] headers) throws IOException
	{
		this(filename, separator, headers, false);
	}
	/**
	 * Create a new writer instance with the specified filename, header values,
	 * whether to keep the number of columns consistent, and the default separator
	 * @param filename The filename of the file to be written
	 * @param headers The values of each field in the header row
	 * @param keepNumberOfColumnsConsistent If true, each line that is written will have the same number of fields, which is based
	 * off the number of fields in the header row (which must be supplied if this is true).
	 * If a line that is to be written has fewer fields, then empty strings are added
	 * on to the end of the row.  If a line that is to be written has more fields, then
	 * the extra fields are ignored.
	 * @throws IOException If the file cannot be created or written
	 * @throws NullPointerException If filename is null
	 * @throws NullPointerException If headers is null
	 */
	public CSVWriter(String filename, String[] headers, boolean keepNumberOfColumnsConsistent) throws IOException
	{
		this(filename, DEFAULT_SEPARATOR, headers, keepNumberOfColumnsConsistent);
	}
	/**
	 * Create a new writer instance with the specified filename, separator, header values,
	 * and whether to keep the number of columns consistent.
	 * @param filename The filename of the file to be written
	 * @param separator The character to be used as the field separator
	 * @param headers The values of each field in the header row
	 * @param keepNumberOfColumnsConsistent If true, each line that is written will have the same number of fields, which is based
	 * off the number of fields in the header row (which must be supplied if this is true).
	 * If a line that is to be written has fewer fields, then empty strings are added
	 * on to the end of the row.  If a line that is to be written has more fields, then
	 * the extra fields are ignored.
	 * @throws IOException If the file cannot be created or written
	 * @throws NullPointerException If filename is null
	 * @throws NullPointerException If headers is null
	 */
	public CSVWriter(String filename, char separator, String[] headers, boolean keepNumberOfColumnsConsistent) throws IOException
	{
		if (headers == null)
			throw new NullPointerException("Headings array is null.");
		
		constructorCommon(filename, separator);
		this.keepNumberOfColumnsConsistent = keepNumberOfColumnsConsistent;
		if (keepNumberOfColumnsConsistent)
			numCols = headers.length;
		writeHeaders(headers);
	}
	
	/**
	 * Construct a new CSVWriter with the specified filename, default separator, and
	 * whether to keep the number of columns consistent based off the specified ResultSet.
	 * This constructor will automatically write the header line based off of the ResultSet
	 * column headers.  If specified by the writeImmediately value, it will then automatically
	 * write all of the rows from the ResultSet to the file.
	 * @param filename The filename of the file to be written
	 * @param rs The ResultSet from which to base this CSVWriter
	 * @param keepNumberOfColumnsConsistent If true, each line that is written will have the same number of fields, which is based
	 * off the number of fields in the header row (which must be supplied if this is true).
	 * If a line that is to be written has fewer fields, then empty strings are added
	 * on to the end of the row.  If a line that is to be written has more fields, then
	 * the extra fields are ignored.
	 * @param writeImmediately If true, all rows from this ResultSet will immediately be written to the file.
	 * @throws SQLException if there is an error with the ResuletSet
	 * @throws IOException If the file cannot be created or written
	 */
	public CSVWriter(String filename, ResultSet rs, boolean keepNumberOfColumnsConsistent, boolean writeImmediately) throws SQLException, IOException
	{
		this(filename, DEFAULT_SEPARATOR, rs, keepNumberOfColumnsConsistent, writeImmediately);
	}
	
	/**
	 * Construct a new CSVWriter with the specified filename, separator, and
	 * whether to keep the number of columns consistent based off the specified ResultSet.
	 * This constructor will automatically write the header line based off of the ResultSet
	 * column headers.  If specified by the writeImmediately value, it will then automatically
	 * write all of the rows from the ResultSet to the file.
	 * @param filename The filename of the file to be written
	 * @param separator The character to be used as the field separator
	 * @param rs The ResultSet from which to base this CSVWriter
	 * @param keepNumberOfColumnsConsistent If true, each line that is written will have the same number of fields, which is based
	 * off the number of fields in the header row (which must be supplied if this is true).
	 * If a line that is to be written has fewer fields, then empty strings are added
	 * on to the end of the row.  If a line that is to be written has more fields, then
	 * the extra fields are ignored.
	 * @param writeImmediately If true, all rows from this ResultSet will immediately be written to the file.
	 * @throws SQLException if there is an error with the ResuletSet
	 * @throws IOException If the file cannot be created or written
	 */
	public CSVWriter(String filename, char separator, ResultSet rs, boolean keepNumberOfColumnsConsistent, boolean writeImmediately) throws SQLException, IOException
	{
		ResultSetMetaData md = rs.getMetaData();
		String[] headers = new String[md.getColumnCount()];
		for (int i = 0; i < headers.length; i++)
			headers[i] = md.getCatalogName(i);
		
		constructorCommon(filename, separator);
		this.keepNumberOfColumnsConsistent = keepNumberOfColumnsConsistent;
		if (keepNumberOfColumnsConsistent)
			numCols = headers.length;
		writeHeaders(headers);
		if (writeImmediately)
			writeMultipleRows(rs);
	}
	/**
	 * Performs the common tasks between the two main constructor variants (headers vs no headers)
	 * @param filename The filename of the file to be written
	 * @param separator The character to be used as the field separator
	 * @throws IOException If the file cannot be created or written
	 * @throws NullPointerException If filename is null
	 */
	private void constructorCommon(String filename, char separator) throws IOException
	{
		writer = new BufferedWriter(new FileWriter(new File(filename)));
		this.separator = separator;
		separatorString = Character.toString(separator);
		numColsThisRow = 0;
	}
	/**
	 * Write the header row to the file as the first row
	 * @param headers The column names to write
	 * @throws IOException If the file cannot be written
	 */
	private void writeHeaders(String[] headers) throws IOException
	{
		writeRow(headers);
	}
	
	/**
	 * Write the specified fields to the file as an individual row.
	 * If keepNumberOfColumnsConsistent is true, then fields will
	 * be added or removed according to the rules specified above.
	 * If fields is null, an empty line is written.
	 * @param fields The values to be written
	 * @throws IOException If the file cannot be written
	 * @throws IllegalStateException If a partial row has been written
	 * (using the writeField and writeMultipleFields methods) without a call to endRow
	 */
	public void writeRow(String[] fields) throws IOException
	{
		if (!newRow)
			throw new IllegalStateException("Cannot write an entire row when a partial" +
					"row has already been written.");
		if (fields != null)
		{
			if (keepNumberOfColumnsConsistent)
			{
				int i;
				//Write until the column limit or the end of the array, whichever comes first.
				//This automatically truncates extra fields
				for (i = 0; i < fields.length && i < numCols; i++)
				{
					if (i > 0)
						writer.write(separator);
					writer.write(encode(fields[i]));
				}
				//Write additional blank fields, if necessary
				if (i < numCols)
					while (i++ < numCols)
						writer.write(separator);
			}
			//Write the entire array, no more and no less
			else
				for (int i = 0; i < fields.length; i++)
				{
					if (i > 0)
						writer.write(separator);
					writer.write(encode(fields[i]));
				}
		}
		//If a blank line is to be written, we may still need to add the blank fields
		else if (keepNumberOfColumnsConsistent)
			for (int i = 1; i < numCols; i++)
				writer.write(separator);
		
		writer.newLine();
		numColsThisRow = 0;
		newRow = true;
	}
	/**
	 * Write multiple rows to the file.
	 * @param lines A list of individual rows to write
	 * @throws IOException If the file cannot be written
	 * @throws NullPointerException If lines is null
	 * @throws IllegalStateException If a partial row has been written
	 * (using the writeField and writeMultipleFields methods) without a call to endRow
	 */
	public void writeMultipleRows(List<String[]> lines) throws IOException
	{
		for (String[] arr : lines)
			writeRow(arr);
	}
	/**
	 * Write a single field to the file.  This method can be used
	 * to write a row to the file piecemeal instead of all at once, should
	 * it be necessary.  When using this method, be sure to call the endRow
	 * method to end the row when appropriate, otherwise calls to writeRow and writeMultipleRows
	 * will result in an exception.  If keepNumberOfColumnsConsistent is true and the
	 * maximum number of columns have already been written on this row, then
	 * this method returns false.
	 * @param field The field to be written
	 * @return False if keepNumberOfColumnsConsistent is true and the
	 * maximum number of columns have already been written on this row; true otherwise.
	 * @throws IOException If the file cannot be written
	 */
	public boolean writeField(String field) throws IOException
	{
		//If this row has already had something written to it, append a separator
		if (keepNumberOfColumnsConsistent && numColsThisRow == numCols)
			return false;
		if (!newRow)
			writer.write(separator);
		//Write this field and indicate that we are in the middle of a row
		writer.write(encode(field));
		newRow = false;
		numColsThisRow++;
		return true;
	}
	/**
	 * Write multiple fields to the file, but not an entire row.  This method can be used
	 * to write a row to the file piecemeal instead of all at once, should
	 * it be necessary.  When using this method, be sure to call the endRow
	 * method to end the row when appropriate, otherwise calls to writeRow and writeMultipleRows
	 * will result in an exception. If keepNumberOfColumnsConsistent is true, then this method
	 * will only write the array only up to the column limit (eg, if this row already has
	 * two fields, the number of columns is four, and the fields array has three elements,
	 * then the first two will be written) and false is returned.
	 * @param fields The fields to be written
	 * @return false if keepNumberOfColumnsConsistent is true and not all
	 * fields could be written; true otherwise.
	 * @throws IOException If the file cannot be written
	 */
	public boolean writeMultipleFields(String[] fields) throws IOException
	{
		//If this row has already had something written to it, append a separator
		if (keepNumberOfColumnsConsistent && numColsThisRow == numCols)
			return false;
		if (!newRow)
			writer.write(separator);
		//Write these fields and indicate that we are in the middle of a row
		int i;
		for (i = 0; i < fields.length && numColsThisRow++ < numCols; i++)
		{
			if (i > 0)
				writer.write(separator);
			writer.write(encode(fields[i]));
		}
		newRow = false;
		return i == fields.length;
	}
	/**
	 * End a row that has been written to piecemeal using writeField and writeMultipleFields.
	 * This method must be called after the above methods have been used before the writeRow
	 * and writeMultipleRows methods may be called again.  If keepColumnNumbersConsistent is true,
	 * then this method will also pad the row with extra fields if necessary.
	 * @throws IOException If the file cannot be written
	 */
	public void endRow() throws IOException
	{
		while (keepNumberOfColumnsConsistent && numColsThisRow++ < numCols)
			writer.write(separatorString);
		writer.newLine();
		newRow = true;
		numColsThisRow = 0;
	}
	/**
	 * Write a blank row in the file, including empty fields if keepColumnNumbersConsistent is true.
	 * This method behaves exactly the same as endRow, but is provided for code clarity, as
	 * endRow serves a more specific purpose (ie, ending a row that was written piecemeal).
	 * @throws IOException If the file cannot be written
	 */
	public void writeEmptyRow() throws IOException
	{
		endRow();
	}
	
	/**
	 * Write the current row of the ResultSet to the file.  At least one call to rs.next()
	 * must occur before this method is called.
	 * @param rs The ResultSet containing the data to write
	 * @throws SQLException if there is an error in the ResultSet
	 * @throws IOException if there is an error writing the file
	 */
	public void writeRow(ResultSet rs) throws SQLException, IOException
	{
		ResultSetMetaData rsmd = rs.getMetaData();
		for (int i = 0; i < rsmd.getColumnCount(); i++)
			writeField(rs.getString(i));
		endRow();
	}
	/**
	 * Write all remaining rows from this ResultSet to the file.  This method will first
	 * call rs.next(), and then write each row.  Thus, to write all data from a ResultSet
	 * using this method, rs.next() should not be called before calling this method.
	 * @param rs The ResultSet containing the data to write
	 * @throws SQLException if there is an error in the ResultSet
	 * @throws IOException if there is an error writing the file
	 */
	public void writeMultipleRows(ResultSet rs) throws SQLException, IOException
	{
		while (rs.next())
			writeRow(rs);
	}
	
	@Override
	public void close() throws IOException
	{
		writer.close();
	}
	/**
	 * Encode the specified String using this instance's separator.
	 * @param s The String to be encoded
	 * @return The encoded String
	 */
	public String encode(String s)
	{
		boolean containsQuote = s.contains ("\"");
		boolean needsEscape = containsQuote || s.contains(separatorString) || s.contains("\n");
		if (needsEscape)
		{
			if (containsQuote)
				s = s.replace("\"", "\"\"");
			s = "\"" + s + "\"";
		}
		return s;
	}
}
